import numpy as np
import matplotlib.pyplot as plt
import ast
filename="nn_data/data2.txt"
data=[]
raw_data=[]
with open(filename) as file:
    for line in file:
        a_data=line.replace(',',' ').split()
        score1=ast.literal_eval(a_data[0])
        score2=ast.literal_eval(a_data[1])
        label=int(a_data[2])
        data.append([score1,score2,label])
      
    data=np.array(data)  
#决策边界
def boundary(x1,w1,w2):
    return (1*w2[2]-(w1[2][0]*w2[0]+w1[2][1]*w2[1]*1)-(w1[0][0]*w2[0]+w1[1][0]*w2[1])*x1)/(w1[0][1]*w2[0]+w1[1][1]*w2[1])
#激活函数
def sigmoid(x):
    return 1/(1+np.exp(-x))
#预测
def prediction(x):
    if x>=0.5:
        return 1
    else:
        return 0
#神经网络类    
class nn:
    def __init__(self,data,batch_size,lr,train_count):
        self.data=data
        self.batch_size=batch_size
        self.lr=lr
        self.train_count=train_count
        #初始设为0
        self.w1=np.zeros((3,2))
        self.w2=np.zeros(3)
        
    def grad1(self,out):
        return out*(1-out)
    
    def grad2(self,out,label):
        return 2*(out-label)*out*(1-out)
    
    def calculate(self,data):
        m=len(data)
        loss=0
        acc=[]
        dw1=np.zeros((3,2))
        dw2=np.zeros(3)
        for i in range(len(data)):
            x1,x2,label=data[i]
            x_=np.array([x1,x2,1])
            z1=sigmoid(self.w1.T@x_)
            z_=np.array([z1[0],z1[1],1])
            z2=sigmoid(np.dot(z_,self.w2))
            loss+=((label-z2)**2)/m
            d2=self.grad2(z2,label)
            dw2+=d2*z_
            for j in range(3):
                dw1[j]+=d2*self.w2[j]*z1

            acc.append(prediction(z2)==label)

        self.w1-=self.lr*dw1
        self.w2-=self.lr*dw2
        return acc,loss
            
    def train(self):
        Accuracy=[]
        Loss=[]
        for i in range(self.train_count):
            np.random.shuffle(self.data)
            for j in range(len(self.data)//self.batch_size):
                train_data=self.data[j*self.batch_size:(j+1)*self.batch_size:1]
                acc,loss=self.calculate(train_data)
                accuracy=acc.count(1)*100/len(train_data)
                if j==0:
                    Loss.append(loss)
                    Accuracy.append(accuracy)
                    print(f"The {i+1} epoch: Acc is {accuracy}% Loss is {loss}")
    
        return Accuracy,Loss,self.w1,self.w2
    
if __name__=="__main__" :
    test=nn(data,50,0.01,5000)
    Accuracy,Loss,w1,w2=test.train()
    # 创建画布和子图
    print("Last w1 is:\n",w1)
    print("Last w2 is:\n",w2)
    fig,ax=plt.subplots(1,3,figsize=(12,4))
    s1=np.linspace(-1,1,200)
    s2=boundary(s1,w1,w2)
    ax[0].plot([i+1 for i in range(5000)],Accuracy)
    ax[0].set_xlabel("Epoch")
    ax[0].set_ylabel("Acc")
    ax[1].plot([i+1 for i in range(5000)],Loss)
    ax[1].set_xlabel("Epoch")
    ax[1].set_ylabel("Loss")
    labels=data[:, 2]
    colors=['r' if k==1 else 'g' for k in labels]
    ax[2].scatter(data[:, 0],data[:, 1],c=colors)
    ax[2].set_xlabel("x1")
    ax[2].set_ylabel("x2")
    ax[2].plot(s1,s2,c='b')
    # 调整子图间的间距
    plt.tight_layout()
    plt.show()
    
